"""
速率限制中间件模块
提供API访问频率限制功能
"""
import time
from functools import wraps
from flask import request, current_app, g

from api.utils.response_utils import rate_limit_response

def rate_limit(limit=10, per=60):
    """
    API速率限制装饰器
    
    Args:
        limit: 时间窗口内允许的最大请求数
        per: 时间窗口大小（秒）
        
    Returns:
        装饰后的函数
    """
    def decorator(f):
        # 使用字典存储请求计数
        # 注意：这只是简单的内存存储，重启后会重置
        # 生产环境应该使用Redis或其他持久化存储
        # 将limits字典作为闭包变量，而不是直接使用current_app
        # 这样可以避免在模块导入时访问current_app
        limits = {}
        
        @wraps(f)
        def decorated(*args, **kwargs):
            # 在请求处理时获取current_app
            # 这样可以确保在Flask应用上下文中访问
            nonlocal limits
            
            if not hasattr(current_app, '_rate_limits'):
                current_app._rate_limits = limits
            else:
                # 使用应用上下文中的limits
                limits = current_app._rate_limits
            
            # 获取客户端IP地址
            client_ip = get_client_ip()
            
            # 生成速率限制键
            limit_key = f"{client_ip}:{f.__name__}"
            
            # 获取当前时间戳
            current_time = time.time()
            
            # 初始化或清理请求记录
            if limit_key not in limits:
                limits[limit_key] = []
            
            # 清理过期的请求记录
            limits[limit_key] = [t for t in limits[limit_key] if current_time - t < per]
            
            # 检查是否超过限制
            if len(limits[limit_key]) >= limit:
                # 计算还需要等待多久才能再次请求
                wait_time = per - (current_time - limits[limit_key][0])
                
                # 构建响应头
                headers = {
                    'X-RateLimit-Limit': str(limit),
                    'X-RateLimit-Remaining': '0',
                    'X-RateLimit-Reset': str(int(current_time + wait_time))
                }
                
                # 返回速率限制响应
                response = rate_limit_response(
                    f'Rate limit exceeded. Try again in {int(wait_time)} seconds.'
                )
                
                # 添加速率限制头
                for key, value in headers.items():
                    response[0].headers[key] = value
                
                return response
            
            # 记录请求
            limits[limit_key].append(current_time)
            
            # 计算剩余请求数
            remaining = limit - len(limits[limit_key])
            
            # 构建响应头
            headers = {
                'X-RateLimit-Limit': str(limit),
                'X-RateLimit-Remaining': str(remaining),
                'X-RateLimit-Reset': str(int(current_time + per))
            }
            
            # 执行原始函数
            response = f(*args, **kwargs)
            
            # 如果响应是元组，分别处理
            if isinstance(response, tuple):
                # 检查是否是Flask响应对象
                if hasattr(response[0], 'headers'):
                    for key, value in headers.items():
                        response[0].headers[key] = value
            elif hasattr(response, 'headers'):
                # 直接是响应对象
                for key, value in headers.items():
                    response.headers[key] = value
            
            return response
        
        return decorated
    
    return decorator

def get_client_ip():
    """
    获取客户端IP地址
    
    Returns:
        str: 客户端IP地址
    """
    # 检查代理头
    for header in ['X-Forwarded-For', 'X-Real-IP', 'Client-IP']:
        ip = request.headers.get(header)
        if ip:
            # 如果X-Forwarded-For有多个IP，取第一个
            if header == 'X-Forwarded-For':
                ip = ip.split(',')[0].strip()
            return ip
    
    # 如果没有代理头，使用REMOTE_ADDR
    return request.remote_addr or '127.0.0.1'

class RateLimiter:
    """
    速率限制器类
    提供更灵活的速率限制功能
    """
    def __init__(self):
        self.limits = {}
        self.global_limits = {}
    
    def check_limit(self, key, limit, per):
        """
        检查是否超过限制
        
        Args:
            key: 限制键
            limit: 时间窗口内允许的最大请求数
            per: 时间窗口大小（秒）
            
        Returns:
            tuple: (是否允许, 等待时间)
        """
        current_time = time.time()
        
        if key not in self.limits:
            self.limits[key] = []
        
        # 清理过期的请求记录
        self.limits[key] = [t for t in self.limits[key] if current_time - t < per]
        
        if len(self.limits[key]) >= limit:
            wait_time = per - (current_time - self.limits[key][0])
            return False, wait_time
        
        return True, 0
    
    def record_request(self, key):
        """
        记录请求
        
        Args:
            key: 限制键
        """
        if key not in self.limits:
            self.limits[key] = []
        
        self.limits[key].append(time.time())
    
    def get_remaining(self, key, limit, per):
        """
        获取剩余请求数
        
        Args:
            key: 限制键
            limit: 时间窗口内允许的最大请求数
            per: 时间窗口大小（秒）
            
        Returns:
            int: 剩余请求数
        """
        current_time = time.time()
        
        if key not in self.limits:
            return limit
        
        # 清理过期的请求记录
        self.limits[key] = [t for t in self.limits[key] if current_time - t < per]
        
        return max(0, limit - len(self.limits[key]))
    
    def get_reset_time(self, key, per):
        """
        获取限制重置时间
        
        Args:
            key: 限制键
            per: 时间窗口大小（秒）
            
        Returns:
            float: 重置时间戳
        """
        if key not in self.limits or not self.limits[key]:
            return time.time()
        
        return self.limits[key][0] + per

# 创建全局速率限制器实例
global_rate_limiter = RateLimiter()

def global_rate_limit(limit=100, per=60):
    """
    全局速率限制装饰器
    使用全局速率限制器实例
    
    Args:
        limit: 时间窗口内允许的最大请求数
        per: 时间窗口大小（秒）
        
    Returns:
        装饰后的函数
    """
    def decorator(f):
        @wraps(f)
        def decorated(*args, **kwargs):
            # 获取客户端IP地址
            client_ip = get_client_ip()
            
            # 生成限制键
            limit_key = f"{client_ip}:global"
            
            # 检查限制
            allowed, wait_time = global_rate_limiter.check_limit(limit_key, limit, per)
            
            if not allowed:
                # 构建响应头
                headers = {
                    'X-RateLimit-Limit': str(limit),
                    'X-RateLimit-Remaining': '0',
                    'X-RateLimit-Reset': str(int(time.time() + wait_time))
                }
                
                # 返回速率限制响应
                response = rate_limit_response(
                    f'Global rate limit exceeded. Try again in {int(wait_time)} seconds.'
                )
                
                # 添加响应头
                for key, value in headers.items():
                    response[0].headers[key] = value
                
                return response
            
            # 记录请求
            global_rate_limiter.record_request(limit_key)
            
            # 计算剩余请求数
            remaining = global_rate_limiter.get_remaining(limit_key, limit, per)
            reset_time = global_rate_limiter.get_reset_time(limit_key, per)
            
            # 构建响应头
            headers = {
                'X-RateLimit-Limit': str(limit),
                'X-RateLimit-Remaining': str(remaining),
                'X-RateLimit-Reset': str(int(reset_time))
            }
            
            # 执行原始函数
            response = f(*args, **kwargs)
            
            # 添加响应头
            if isinstance(response, tuple) and hasattr(response[0], 'headers'):
                for key, value in headers.items():
                    response[0].headers[key] = value
            elif hasattr(response, 'headers'):
                for key, value in headers.items():
                    response.headers[key] = value
            
            return response
        
        return decorated
    
    return decorator

def sliding_window_rate_limit(limit=10, per=60):
    """
    滑动窗口速率限制装饰器
    更精确的速率限制实现
    
    Args:
        limit: 时间窗口内允许的最大请求数
        per: 时间窗口大小（秒）
        
    Returns:
        装饰后的函数
    """
    def decorator(f):
        # 使用字典存储请求时间戳
        if not hasattr(current_app, '_sliding_limits'):
            current_app._sliding_limits = {}
        
        sliding_limits = current_app._sliding_limits
        
        @wraps(f)
        def decorated(*args, **kwargs):
            # 获取客户端IP地址
            client_ip = get_client_ip()
            
            # 生成限制键
            limit_key = f"{client_ip}:{f.__name__}"
            
            # 获取当前时间戳
            current_time = time.time()
            
            # 初始化
            if limit_key not in sliding_limits:
                sliding_limits[limit_key] = []
            
            # 滑动窗口：移除时间窗口外的请求
            window_start = current_time - per
            sliding_limits[limit_key] = [t for t in sliding_limits[limit_key] if t > window_start]
            
            # 检查是否超过限制
            if len(sliding_limits[limit_key]) >= limit:
                # 计算最早的请求时间，确定重置时间
                earliest_request = min(sliding_limits[limit_key])
                reset_time = earliest_request + per
                wait_time = reset_time - current_time
                
                headers = {
                    'X-RateLimit-Limit': str(limit),
                    'X-RateLimit-Remaining': '0',
                    'X-RateLimit-Reset': str(int(reset_time))
                }
                
                response = rate_limit_response(
                    f'Rate limit exceeded. Try again in {int(wait_time)} seconds.'
                )
                
                for key, value in headers.items():
                    response[0].headers[key] = value
                
                return response
            
            # 记录请求
            sliding_limits[limit_key].append(current_time)
            
            # 计算剩余请求数
            remaining = limit - len(sliding_limits[limit_key])
            
            # 计算重置时间（基于最早的请求）
            if sliding_limits[limit_key]:
                earliest_request = min(sliding_limits[limit_key])
                reset_time = earliest_request + per
            else:
                reset_time = current_time + per
            
            headers = {
                'X-RateLimit-Limit': str(limit),
                'X-RateLimit-Remaining': str(remaining),
                'X-RateLimit-Reset': str(int(reset_time))
            }
            
            # 执行原始函数
            response = f(*args, **kwargs)
            
            # 添加响应头
            if isinstance(response, tuple) and hasattr(response[0], 'headers'):
                for key, value in headers.items():
                    response[0].headers[key] = value
            elif hasattr(response, 'headers'):
                for key, value in headers.items():
                    response.headers[key] = value
            
            return response
        
        return decorated
    
    return decorator

# 导出
__all__ = [
    'rate_limit',
    'global_rate_limit',
    'sliding_window_rate_limit',
    'RateLimiter',
    'global_rate_limiter',
    'get_client_ip'
]